#!/bin/env python
# -*- coding: utf-8 -*-
"""
内核编译缓存管理器

提供Taichi内核编译结果的缓存管理，优化重复编译的性能问题。
"""

import os
import time
import hashlib
import taichi as ti
import json
import logging
from typing import Dict, Optional, Any, Tuple

class KernelCacheManager:
    """
    Taichi内核编译缓存管理类
    
    利用Taichi的离线缓存（offline cache）机制，管理内核编译结果，
    避免重复编译带来的性能损耗。
    """
    
    def __init__(self, cache_dir: Optional[str] = None, debug: bool = False):
        """
        初始化内核缓存管理器
        
        :param cache_dir: 缓存目录，默认为 ~/.taichi_kernel_cache
        :param debug: 是否启用调试模式
        """
        self.debug = debug
        self.logger = self._setup_logger()
        
        # 设置缓存目录
        if cache_dir is None:
            home_dir = os.path.expanduser("~")
            self.cache_dir = os.path.join(home_dir, ".taichi_kernel_cache")
        else:
            self.cache_dir = os.path.expanduser(cache_dir)
            
        # 创建缓存目录（如果不存在）
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir, exist_ok=True)
            
        # 缓存索引文件路径
        self.index_file = os.path.join(self.cache_dir, "cache_index.json")
        
        # 加载索引
        self.cache_index = self._load_index()
        
        self.logger.info(f"内核缓存管理器初始化完成，使用缓存目录: {self.cache_dir}")
    
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger("KernelCacheManager")
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            
        logger.setLevel(logging.DEBUG if self.debug else logging.INFO)
        return logger
    
    def _load_index(self) -> Dict:
        """加载缓存索引"""
        if os.path.exists(self.index_file):
            try:
                with open(self.index_file, 'r') as f:
                    return json.load(f)
            except (json.JSONDecodeError, FileNotFoundError):
                self.logger.warning(f"无法读取索引文件，将创建新的索引")
                return {}
        return {}
    
    def _save_index(self):
        """保存缓存索引"""
        try:
            with open(self.index_file, 'w') as f:
                json.dump(self.cache_index, f, ensure_ascii=False, indent=2)
        except Exception as e:
            self.logger.error(f"保存索引文件失败: {e}")
    
    def _generate_cache_key(self, func_name: str, args: Tuple, kwargs: Dict) -> str:
        """
        生成缓存键
        
        :param func_name: 函数名称
        :param args: 位置参数
        :param kwargs: 关键字参数
        :return: 缓存键
        """
        # 构建一个包含函数名和参数的字符串
        key_str = f"{func_name}:{str(args)}:{str(sorted(kwargs.items()))}"
        # 使用MD5生成唯一哈希值
        return hashlib.md5(key_str.encode()).hexdigest()
    
    def initialize_taichi(self, **kwargs):
        """
        初始化Taichi运行环境
        
        包含对离线缓存的设置，确保Taichi使用我们指定的缓存目录。
        
        :param kwargs: 传递给ti.init的参数
        """
        # 设置离线缓存相关参数
        if 'offline_cache' not in kwargs:
            kwargs['offline_cache'] = True
            
        if 'offline_cache_file_path' not in kwargs:
            kwargs['offline_cache_file_path'] = self.cache_dir
            
        # 初始化Taichi
        ti.init(**kwargs)
        self.logger.info(f"Taichi初始化完成，离线缓存设置为: {kwargs.get('offline_cache')}")
    
    def prepare_kernel(self, init_params: Optional[Dict[str, Any]] = None):
        """
        准备Taichi内核环境
        
        :param init_params: 传递给ti.init的参数字典
        """
        params = init_params or {}
        self.initialize_taichi(**params)
    
    def clear_cache(self, days_old: Optional[int] = None):
        """
        清理缓存
        
        :param days_old: 如果设置，将删除超过指定天数的缓存
        """
        if days_old is not None:
            now = time.time()
            max_age = days_old * 86400  # 天数转换为秒
            
            # 过滤出需要保留的缓存项
            new_index = {}
            for key, info in self.cache_index.items():
                if now - info.get('timestamp', 0) < max_age:
                    new_index[key] = info
                else:
                    # 删除过期的缓存文件
                    cache_file = info.get('file_path')
                    if cache_file and os.path.exists(cache_file):
                        try:
                            os.remove(cache_file)
                            self.logger.info(f"已删除过期缓存: {cache_file}")
                        except OSError as e:
                            self.logger.error(f"删除缓存文件失败: {e}")
            
            self.cache_index = new_index
            self._save_index()
            self.logger.info(f"已清理超过 {days_old} 天的缓存")
        else:
            # 清空所有缓存
            for info in self.cache_index.values():
                cache_file = info.get('file_path')
                if cache_file and os.path.exists(cache_file):
                    try:
                        os.remove(cache_file)
                    except OSError as e:
                        self.logger.error(f"删除缓存文件失败: {e}")
                        
            self.cache_index = {}
            self._save_index()
            self.logger.info("已清空所有缓存")
    
    def cache_stats(self) -> Dict:
        """
        获取缓存统计信息
        
        :return: 包含缓存统计的字典
        """
        total_size = 0
        file_count = 0
        
        # 检查每个缓存文件并计算总大小
        for info in self.cache_index.values():
            cache_file = info.get('file_path')
            if cache_file and os.path.exists(cache_file):
                file_count += 1
                total_size += os.path.getsize(cache_file)
        
        # 计算缓存命中率
        hits = sum(1 for info in self.cache_index.values() if info.get('hits', 0) > 0)
        misses = len(self.cache_index) - hits
        hit_rate = hits / (hits + misses) if hits + misses > 0 else 0
        
        return {
            'file_count': file_count,
            'total_size_bytes': total_size,
            'total_size_mb': total_size / (1024 * 1024),
            'hit_count': hits,
            'miss_count': misses,
            'hit_rate': hit_rate
        }
    
    def update_cache_record(self, cache_key: str, file_path: str, hit: bool = False):
        """
        更新缓存记录
        
        :param cache_key: 缓存键
        :param file_path: 缓存文件路径
        :param hit: 是否为缓存命中
        """
        if cache_key in self.cache_index:
            self.cache_index[cache_key]['timestamp'] = time.time()
            if hit:
                self.cache_index[cache_key]['hits'] = self.cache_index[cache_key].get('hits', 0) + 1
        else:
            self.cache_index[cache_key] = {
                'file_path': file_path,
                'timestamp': time.time(),
                'hits': 1 if hit else 0
            }
        
        self._save_index()

# 创建一个全局的内核缓存管理器实例
kernel_cache_manager = KernelCacheManager()

def enable_kernel_cache(arch=None, debug=False, **kwargs):
    """
    启用内核缓存的快捷函数
    
    :param arch: Taichi架构，如ti.cpu, ti.gpu等
    :param debug: 是否开启调试模式
    :param kwargs: 其他传递给ti.init的参数
    :return: 内核缓存管理器实例
    """
    params = {
        'arch': arch or ti.gpu,
        'debug': debug,
        'kernel_profiler': True,
        'offline_cache': True,
        **kwargs
    }
    if arch == ti.cpu:
        return ti.init(arch=ti.cpu, debug=debug, **kwargs)
    else:
        kernel_cache_manager.prepare_kernel(params)
        return kernel_cache_manager 